%% Replication File for Baqaee, Burstein, and Koike-Mori (2023) %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
% Prepared by Yasutaka Koike-Mori (UCLA) 
%
% Description: This function calculates the Money Metric Utility Function using initial-year base prices.
% The order for all matrices is (I,N,T) where:
% I denotes the number of goods
% N denotes the number of households.
% T denotes the number of time periods
%
% Inputs: I_vec, B_vec, pvec
%   I_vec: A (1,N,T) matrix representing income.
%   B_vec: A (I,N,T) matrix of Uncompensated Budget shares.
%   pvec : A (I,1,T) matrix representing price.
%   flagAlgorithm: Set to 0 for iterative, 1 for recursive.
%   Note: Prior to use, sort I_vec and B_vec in ascending order by the second dimension of I_vec.
%         The number of households can vary by period. See N to be the maximum number of households across periods, and set data for missing households in a give period to NaN
% 
% Outputs: U_vec
%   U_vec: A (1,N,T) matrix containing Money Metric. For income levels for which we cannot calculate money-metric, U_vec=NaN.
%
% Internal Functions: CalMoneyMetric_t, compensated_b, updateMM, makefunction
%   CalMoneyMetric_t: Calculates money metric at time t.
%   uncompensated_B : Obtain budget share as a function of income without extrapolation.
%   checkIterative  : Final check for iterative algorithm.
%   makefunction    : Interpolates allowing for NaN values (Allow approximations outside the support).
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function U_vec = CalMoneyMetric(I_vec, B_vec, pvec, flagAlgorithm)
    % Get sizes
    [~, ~, T] = size(B_vec);
    % Initializations
    dlogp = [diff(log(pvec),[],3)];
    B_vec = B_vec./sum(B_vec);
    U_vec(:,:,1) = I_vec(:,:,1); % The money metric for the first time period is assumed to be the income for that period
    F_I_U  = cell(1,T);  F_U_I  = cell(1,T);
    F_B_I = uncompensated_B(B_vec,I_vec,T); %make a budget fucntion of Income without extrapolation of Engel curve 
    if flagAlgorithm == 1 % check if iterative or recursive
        Maxiter = 20;
    else
        Maxiter = 1;
    end

    % Loop over T
    for t = 2:T    
        if mod(t,10) == 0
            disp(['t: ', num2str(t)]);
        end
        % Initialize parameters related iteration.
        tau = 1; err = 100; tol = 10^(-4);
        % Iterate until convergence
        while err > tol
            [U_vec,F_I_U,F_U_I] = CalMoneyMetric_t(t,dlogp, I_vec, B_vec, U_vec, F_B_I,F_I_U,F_U_I,tau);
            if tau > 1
                err = max(abs(log(U_vec(:,:,t)) - logU_vec_old));
            end            
            logU_vec_old = log(U_vec(:,:,t));            
            tau = tau + 1; %iteration number for the recursive algorithm
            if tau > Maxiter
                break
            end
        end
    end 
end

% function to compute money metric at time t
function [U_vec,F_I_U,F_U_I] = CalMoneyMetric_t(t,dlogp, I_vec, B_vec, U_vec, F_B_I, F_I_U, F_U_I, tau)
    % Get sizes
    [I, N, ~] = size(B_vec);
    % Initialization
    B_Istar = zeros(I,N);  %
    Integ = zeros(1,N);
    % Compute utility values for time step t
    if tau ==1
        % Make function of I(U) and U(I). This function allows approximation outside of support
        F_I_U{1,t-1} = makefunction(log(U_vec(:,:,t-1)),log(I_vec(:,:,t-1)));
        F_U_I{1,t-1} = makefunction(log(I_vec(:,:,t-1)),log(U_vec(:,:,t-1))); 
        % Initial Guess:
        logUt = F_U_I{1,t-1}(log(I_vec(:,:,t))); 
    else % If tau > 1, use the value of previous iteration.
        logUt = log(U_vec(:,:,t));  
    end
    % Compute Riemann sum
    for s = 1:t-1
            logIstar_s = F_I_U{1,s}(logUt);
            if s < t-1 
                logIstar_s1 = F_I_U{1,s + 1}(logUt); 
            end
        for i = 1:I % Use trapezoid rule. 
            if s < t-1 
                B_Istar(i,:) = (F_B_I{i,s}(logIstar_s) + F_B_I{i,s+1}(logIstar_s1))/2;
            else 
                B_Istar(i,:) = (F_B_I{i,s}(logIstar_s) + B_vec(i,:,t))/2;
            end
        end 
        B_Istar = B_Istar./sum(B_Istar,1); 
        % Update integral by trapezoid rule
        Integ = Integ + ( - (sum(B_Istar.*dlogp(:,:,s)))  ); 
    end
    % Update Boundary_vec
    % Since F_B_I returns NaN if the income is outside the support of the Engel Curve, U_vec contains only samples without extrapolation of Engel Curve.
    U_vec(:,:,t) = exp(log(I_vec(:,:,t)) + Integ );
    % Final check for iterative solution
    if tau == 1
        U_vec = checkIterative(U_vec,I_vec,F_I_U, t);
    end
end

%% Internal functions
% Final check for iterative solution
function U_vec = checkIterative(U_vec,I_vec,F_I_U, t)
for s = 1:t-1
    maxIs = max(I_vec(:,:,s));
    minIs = min(I_vec(:,:,s));
    Istars = exp(F_I_U{1,s}(log(U_vec(:,:,t))));
    indexs = (Istars >= minIs) & (Istars <= maxIs) ;
    U_vec(:,indexs==0,t) = NaN;
end
end

% Function to calculate conpensated budget share without extrapolation
function f_Bs_I = uncompensated_B(B_vec,I_vec,T)
    for t = 1:T
       for i = 1:size(B_vec,1)
            f_Bs_I{i,t} = interp_I(log(I_vec(:,:,t)),B_vec(i,:,t));
       end
    end
% Function to interpolate Budget shares as a function of I without extrapolation
function F_X = interp_I(X,Y)
    indexToKeep = ~isnan(X) & ~isnan(Y);
    X = X(indexToKeep); Y = Y(indexToKeep);
    [X, indexUN] = unique(X);
    Y = Y(indexUN);
    F_X = griddedInterpolant(X,Y,'linear','none');
end    
end


% Function to interpolate vectors with NaN values
function F_X = makefunction(X,Y)
    indexToKeep = ~isnan(X) & ~isnan(Y);
    X = X(indexToKeep); Y = Y(indexToKeep);    
    [X, indexUN] = unique(X);
    Y = Y(indexUN);
    F_X = griddedInterpolant(X,Y);
end



